{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "initial_id",
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import DataLoader\n",
    "from torchvision import datasets, transforms\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "# 设置随机种子以确保结果可复现\n",
    "torch.manual_seed(42)\n",
    "\n",
    "# 数据预处理\n",
    "transform = transforms.Compose([\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize((0.1307,), (0.3081,))  # MNIST数据集的均值和标准差\n",
    "])\n",
    "\n",
    "# 加载MNIST数据集\n",
    "train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)\n",
    "test_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)\n",
    "\n",
    "# 创建数据加载器\n",
    "batch_size = 64\n",
    "train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)\n",
    "test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)\n",
    "\n",
    "# 查看数据样本\n",
    "def visualize_samples():\n",
    "    samples = next(iter(train_loader))\n",
    "    images, labels = samples\n",
    "\n",
    "    fig, axes = plt.subplots(2, 5, figsize=(12, 5))\n",
    "    axes = axes.flatten()\n",
    "\n",
    "    for i in range(10):\n",
    "        axes[i].imshow(images[i][0], cmap='gray')\n",
    "        axes[i].set_title(f'Label: {labels[i]}')\n",
    "        axes[i].axis('off')\n",
    "\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "\n",
    "# 定义MLP模型\n",
    "class MLP(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(MLP, self).__init__()\n",
    "        # 将28x28的图像展平为784维向量\n",
    "        self.flatten = nn.Flatten()\n",
    "        # 定义多层感知机结构\n",
    "        self.layers = nn.Sequential(\n",
    "            nn.Linear(784, 256),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(256, 128),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(128, 10)  # 输出10个类别（0-9）\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.flatten(x)\n",
    "        x = self.layers(x)\n",
    "        return x\n",
    "\n",
    "# 初始化模型、损失函数和优化器\n",
    "model = MLP()\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = optim.Adam(model.parameters(), lr=0.001)\n",
    "\n",
    "# 训练模型\n",
    "def train(model, train_loader, criterion, optimizer, epochs):\n",
    "    model.train()\n",
    "    train_losses = []\n",
    "    train_acc = []\n",
    "\n",
    "    for epoch in range(epochs):\n",
    "        running_loss = 0.0\n",
    "        correct = 0\n",
    "        total = 0\n",
    "\n",
    "        for inputs, labels in train_loader:\n",
    "            # 梯度清零\n",
    "            optimizer.zero_grad()\n",
    "\n",
    "            # 前向传播\n",
    "            outputs = model(inputs)\n",
    "            loss = criterion(outputs, labels)\n",
    "\n",
    "            # 反向传播和优化\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "\n",
    "            running_loss += loss.item()\n",
    "\n",
    "            # 计算准确率\n",
    "            _, predicted = outputs.max(1)\n",
    "            total += labels.size(0)\n",
    "            correct += predicted.eq(labels).sum().item()\n",
    "\n",
    "        # 计算平均损失和准确率\n",
    "        epoch_loss = running_loss / len(train_loader)\n",
    "        epoch_acc = 100.0 * correct / total\n",
    "\n",
    "        train_losses.append(epoch_loss)\n",
    "        train_acc.append(epoch_acc)\n",
    "\n",
    "        print(f'Epoch {epoch+1}/{epochs}, Loss: {epoch_loss:.4f}, Accuracy: {epoch_acc:.2f}%')\n",
    "\n",
    "    return train_losses, train_acc\n",
    "\n",
    "# 测试模型\n",
    "def test(model, test_loader, criterion):\n",
    "    model.eval()\n",
    "    test_loss = 0\n",
    "    correct = 0\n",
    "    total = 0\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for inputs, labels in test_loader:\n",
    "            outputs = model(inputs)\n",
    "            loss = criterion(outputs, labels)\n",
    "\n",
    "            test_loss += loss.item()\n",
    "            _, predicted = outputs.max(1)\n",
    "            total += labels.size(0)\n",
    "            correct += predicted.eq(labels).sum().item()\n",
    "\n",
    "    # 计算平均损失和准确率\n",
    "    test_loss = test_loss / len(test_loader)\n",
    "    test_acc = 100.0 * correct / total\n",
    "\n",
    "    print(f'Test Loss: {test_loss:.4f}, Test Accuracy: {test_acc:.2f}%')\n",
    "\n",
    "    return test_loss, test_acc\n",
    "\n",
    "# 绘制训练过程\n",
    "def plot_training(train_losses, train_acc):\n",
    "    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 5))\n",
    "\n",
    "    # 绘制损失曲线\n",
    "    ax1.plot(train_losses)\n",
    "    ax1.set_title('Training Loss')\n",
    "    ax1.set_xlabel('Epoch')\n",
    "    ax1.set_ylabel('Loss')\n",
    "\n",
    "    # 绘制准确率曲线\n",
    "    ax2.plot(train_acc)\n",
    "    ax2.set_title('Training Accuracy')\n",
    "    ax2.set_xlabel('Epoch')\n",
    "    ax2.set_ylabel('Accuracy (%)')\n",
    "\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "\n",
    "# 主函数\n",
    "def main():\n",
    "    # 可视化样本\n",
    "    visualize_samples()\n",
    "\n",
    "    # 训练模型\n",
    "    epochs = 5\n",
    "    train_losses, train_acc = train(model, train_loader, criterion, optimizer, epochs)\n",
    "\n",
    "    # 测试模型\n",
    "    test_loss, test_acc = test(model, test_loader, criterion)\n",
    "\n",
    "    # 绘制训练过程\n",
    "    plot_training(train_losses, train_acc)\n",
    "\n",
    "    # 可视化一些预测结果\n",
    "    visualize_predictions()\n",
    "\n",
    "# 可视化预测结果\n",
    "def visualize_predictions():\n",
    "    model.eval()\n",
    "    samples = next(iter(test_loader))\n",
    "    images, labels = samples\n",
    "\n",
    "    with torch.no_grad():\n",
    "        outputs = model(images)\n",
    "        _, predictions = torch.max(outputs, 1)\n",
    "\n",
    "    fig, axes = plt.subplots(2, 5, figsize=(12, 5))\n",
    "    axes = axes.flatten()\n",
    "\n",
    "    for i in range(10):\n",
    "        axes[i].imshow(images[i][0], cmap='gray')\n",
    "        axes[i].set_title(f'真实: {labels[i]}\\n预测: {predictions[i]}')\n",
    "        axes[i].axis('off')\n",
    "\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    main()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
